package com.gallop.mq.core.template;

import com.gallop.mq.MQProperties;
import com.gallop.mq.MQTemplate;
import com.gallop.mq.core.MQExecutorProvider;
import com.gallop.mq.core.MQMessageCodec;
import com.gallop.mq.core.MQMessageDispatcher;
import com.gallop.mq.core.MQReliableRegistry;
import com.gallop.mq.core.dispatcher.DispatchContext;
import com.gallop.mq.core.exception.MQException;
import com.gallop.mq.core.template.handler.registry.MessageDeliveryHandlerRegistry;
import io.nats.client.*;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;

import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.stream.Collectors;

@Slf4j
@Getter
@Setter
public class NatsMQTemplate implements MQTemplate {

    public static final String NATS_SCHEME = "nats://";

    private MQMessageCodec codec;
    private Connection natsConnection;
    private MQMessageDispatcher dispatcher;
    private final AtomicBoolean working = new AtomicBoolean();
    private MQProperties properties;

    private Map<String, MQReliableRegistry.ReliableDeliveryConfig> reliableConfigMap;

    @Override
    public void init(MQProperties properties, MQMessageDispatcher dispatcher, MQMessageCodec codec, MQExecutorProvider executorProvider, MessageDeliveryHandlerRegistry handlerRegistry) throws MQException {
        this.dispatcher = dispatcher;
        this.codec = codec;
        this.properties = properties;
        this.reliableConfigMap = dispatcher.getReliableRegistry().register().stream()
                .collect(Collectors.toMap(MQReliableRegistry.ReliableDeliveryConfig::getTopic, c -> c));
        Options.Builder builder = new Options.Builder();
        if (StringUtils.hasLength(properties.getUser()) && StringUtils.hasLength(properties.getSecret())) {
            builder.userInfo(properties.getUser().toCharArray(), properties.getSecret().toCharArray());
        }
        ExecutorService executor = executorProvider != null ? executorProvider.create() : null;
        if (executor != null) {
            builder.executor(executor);
        }
        Options options = builder.servers(appendSchemeIfAbsent(properties.getHosts()).toArray(new String[0]))
                .connectionListener((conn, type) -> {
                    if (ConnectionListener.Events.DISCONNECTED.equals(type) || ConnectionListener.Events.CLOSED.equals(type)) {
                        this.working.compareAndSet(true, false);
                        log.warn("MQ Nats disconnected.");
                    }
                    if (ConnectionListener.Events.CONNECTED.equals(type)) {
                        this.working.compareAndSet(false, true);
                        log.warn("MQ Nats connected.");
                    }
                    if (ConnectionListener.Events.RECONNECTED.equals(type)) {
                        this.working.compareAndSet(false, true);
                        log.warn("MQ Nats reconnected.");
                    }
                })
                .maxReconnects(-1)
                .reconnectWait(Duration.ofMillis(properties.getReconnectInterval()))
                .maxMessagesInOutgoingQueue(properties.getNats().getMaxMessagesInOutgoingQueueSize())
                .discardMessagesWhenOutgoingQueueFull()
                .build();
        try {
            this.natsConnection = Nats.connect(options);
            Dispatcher natsAsyncDispatcher = this.natsConnection.createDispatcher();

            this.dispatcher.getAllContainers().forEach(container -> {
                String topic = container.getTopic();
                String group = container.getGroup();
                if (StringUtils.hasLength(group)) {
                    natsAsyncDispatcher.subscribe(topic, group, getMessageHandler());
                } else {
                    natsAsyncDispatcher.subscribe(topic, getMessageHandler());
                }
            });
            this.working.set(true);
        } catch (Exception e) {
            throw new MQException("init mq template failed. property=" + properties, e);
        }
    }

    private List<String> appendSchemeIfAbsent(List<String> hosts) {
        return hosts.stream().filter(StringUtils::hasLength)
                .map(host -> host.startsWith(NATS_SCHEME) ? host : NATS_SCHEME + host)
                .collect(Collectors.toList());
    }

    @Override
    public void onDispose() {
        if (this.natsConnection != null) {
            try {
                this.natsConnection.close();
            } catch (InterruptedException e) {
                log.error("", e);
            }
        }
    }

    private MessageHandler getMessageHandler() {
        return msg -> {
            Object message = this.codec.decode(msg.getData());
            if (message == null) {
                log.warn("null MQ message will be ignored.");
                return;
            }
            String queueName = msg.getSubscription().getQueueName();
            this.dispatcher.dispatch(DispatchContext.builder()
                    .topic(msg.getSubject())
                    .group(queueName)
                    .replyTo(msg.getReplyTo())
                    .correlationId(msg.getSID())
                    .message(message)
                    .rpcResponse((correlationId, replyTo, rpcResponse) -> {
                        if (StringUtils.hasLength(replyTo)) {
                            this.send(rpcResponse, replyTo);
                        }
                    }).build());
        };
    }

    @Override
    public MQMessageCodec getCodec() {
        return codec;
    }

    @Override
    public void setCodec(MQMessageCodec codec) {
        this.codec = codec;
    }

    @Override
    public MQMessageDispatcher getDispatcher() {
        return this.dispatcher;
    }

    @Override
    public void setDispatcher(MQMessageDispatcher dispatcher) {
        this.dispatcher = dispatcher;
    }

    @Override
    public void send(Object message, String topic) {
//        if (!this.working.get()) {
//            return;
//        }
        if (!handleReliableSend(topic, message)) {
            log.error("handleReliableSend failed.topic=[{}],message=[{}]", topic, message);
            return;
        }
        this.natsConnection.publish(topic, this.codec.encode(message));
    }

    private boolean handleReliableSend(String topic, Object message) {
        MQReliableRegistry.ReliableDeliveryConfig reliableCfg = this.reliableConfigMap.get(topic);
        if (reliableCfg == null) {
            return true;
        }
        if (reliableCfg.getMessageType().isAssignableFrom(message.getClass())) {
            return reliableCfg.getBeforeSend().apply(reliableCfg, topic, message, null);
        }
        return true;
    }

    @Override
    public <P> void request(Object message, String topic, Duration requestTimeout, Consumer<P> onResponse) {
        this.request(message, topic, requestTimeout, onResponse, (req, err) -> log.error("", err));
    }

    @Override
    @SuppressWarnings("unchecked")
    public <P> void request(Object message, String topic, Duration requestTimeout, Consumer<P> onResponse, BiConsumer<Object, Exception> onError) {
        try {
            Message responseMessage;
            if (requestTimeout == null && properties.getRequestTimeout() > 0) {
                requestTimeout = Duration.ofMillis(properties.getRequestTimeout());
            }
            if (requestTimeout == null) {
                CompletableFuture<Message> future = this.natsConnection.request(topic, this.codec.encode(message));
                responseMessage = future.get();
            } else {
                responseMessage = this.natsConnection.request(topic, this.codec.encode(message), requestTimeout);
            }
            if (responseMessage == null) {
                throw new TimeoutException("async request timeout. message={}" + message + ",topic=" + topic);
            }
            P responseData = (P) this.codec.decode(responseMessage.getData());
            onResponse.accept(responseData);
        } catch (Exception e) {
            onError.accept(message, e);
        }
    }

    @Override
    public <P> P request(Object message, String topic) throws InterruptedException {
        return request(message, topic, properties.getRequestTimeout() <= 0 ? null : Duration.ofMillis(properties.getRequestTimeout()));
    }

    @Override
    @SuppressWarnings("unchecked")
    public <P> P request(Object message, String topic, Duration receiveTimeout) throws InterruptedException {
        Message responseMessage;
        if (receiveTimeout == null) {
            try {
                CompletableFuture<Message> future = this.natsConnection.request(topic, this.codec.encode(message));
                responseMessage = future.get();
            } catch (Exception e) {
                log.error("", e);
                return null;
            }
        } else {
            responseMessage = this.natsConnection.request(topic, this.codec.encode(message), receiveTimeout);
        }
        if (responseMessage == null) {
            return null;
        }
        return (P) this.codec.decode(responseMessage.getData());
    }
}
